#include "ImageProcessor.h"
#include <iostream>
#include <vector>
#include <map>

using namespace std;

bool ImageProcessor::getPedalLowered(double rat)
{
	if(rat < (raisedRatio + loweredRatio) / 2.0)
		return true;
	else
		return false;
}

/* Constructor
   Initializes the image processor with a given image.
   The image is not actually processed, and only its
   statistics are gathered.  This is necessary to
   initialize the dynamically-allocated fields before
   any processing begins.
*/
ImageProcessor::ImageProcessor(IplImage* sample)
{
	width = sample->width;
	height = sample->height;
	step = sample->widthStep;
	channels = sample->nChannels;
	blobField = new int*[height];
	countField = new int*[height];

	upperLeft.setPosition(0, 0);
	upperRight.setPosition(0, 0);
	lowerLeft.setPosition(0, 0);
	lowerRight.setPosition(0, 0);
	
	for(int i = 0; i < height; i++)
	{
		blobField[i] = new int[width];
		countField[i] = new int[width];
	}

}

/*  calibrateStep
    Performs a given step of calibration.
*/
void ImageProcessor::calibrateStep(int step)
{
	switch(step)
	{
		// calibrate white-balance
		/*case 0:
		{
			int midX = width / 2;
			int midY = height / 2;

			conversionR = conversionG = conversionB = 0.0;
			pixelCount = 0;

			for(int i = midY - 10; i <= midY + 10; i++)
			{
				for(int j = midX - 10; j <= midX + 10; j++)
				{
					// access colors
					uchar* red = ((uchar*)(data + step*i + channels*j + 2));
					uchar* green = ((uchar*)(data + step*i + channels*j + 1));
					uchar* blue = ((uchar*)(data + step*i + channels*j));

					// if a dark line, skip
					if(*red + *green + *blue < 40)
						continue;

					pixelCount++;
					// assume paper is 240, 240, 240
					//conversionR += (140 - *red);
					//conversionG += (140 - *green);
					//conversionB += (140 - *blue);
				}
			}

			conversionR /= pixelCount;
			conversionB /= pixelCount;
			conversionG /= pixelCount;

			break;
		}*/


		// calibrate upper left corner
		case 0:
		{
			int x = redBlob.getCOMX();
			int y = redBlob.getCOMY();

			lowerLeft.setPosition(x, y);
			cout << x << " " << y << endl;
			break;
		}

		// calibrate upper right corner
		case 1:
		{
			int x = redBlob.getCOMX();
			int y = redBlob.getCOMY();

			upperRight.setPosition(x, y);
			cout << x << " " << y << endl;
			break;
		}

		// calibrate lower right corner
		case 2:
		{
			int x = redBlob.getCOMX();
			int y = redBlob.getCOMY();

			lowerRight.setPosition(x, y);
			cout << x << " " << y << endl;
			break;
		}
		
		// calibrate upper left corner
		case 3:
		{
			int x = redBlob.getCOMX();
			int y = redBlob.getCOMY();

			upperLeft.setPosition(x, y);
			cout << x << " " << y << endl;
			break;
		}

		// calibrate lowered ratio
		case 4:
		{
			loweredRatio = pedalBlob.getRatio();
			break;
		}

		// calibrate raised ratio
		case 5:
		{
			raisedRatio = pedalBlob.getRatio();
			break;
		}

		// shouldn't happen
		default:
		{
			cout << "Calibration is already complete!" << endl;
			break;
		}
	}
}

ImageProcessor::~ImageProcessor(void)
{
	// cleanup
	for(int i = 0; i < height; i++)
	{
		delete [] blobField[i];
		delete [] countField[i];
	}
	delete [] blobField;
	delete [] countField;

}

/*  processImage
	Processes the next frame.
*/
void ImageProcessor::processImage(IplImage* image)
{
	vector<int> blobs[COLOR_COUNT];
	map<int, int> pixelCount[COLOR_COUNT];
	data = (uchar*)image->imageData;
	int index[COLOR_COUNT];
	redBlob.clear();
	greenBlob.clear();
	blueBlob.clear();
	pedalBlob.clear();

	for(int i = 0; i < COLOR_COUNT; i++)
		index[i] = 0;

	// first assign blobs
	for(int i = 0; i < height; i++)
	{
		for(int j = 0; j < width; j++)
		{
			// clear the fields
			blobField[i][j] = 0;
			countField[i][j] = 0;

			// access colors
			uchar* red = ((uchar*)(data + step*i + channels*j + 2));
			uchar* green = ((uchar*)(data + step*i + channels*j + 1));
			uchar* blue = ((uchar*)(data + step*i + channels*j));

			double total = double(*red) + double(*blue) + double(*green);

			double r = (*red / total);

			double g = (*green / total);
			double b = (*blue / total);

			// if the pixel is very dark, skip over it to prevent noise error
			if (total < 30.0)
			{
				continue;
			}
			// determine if this is skin based on red and green
			else if(r > 0.4 && g > 0.1 && r > g && g > b)
			{
				// can only actually be skin if it is to the left of the keys
				if (j < (upperLeft.getX() + lowerLeft.getX()) / 2)
					blobField[i][j] = SKIN;
			}
			// identify red pixels
			else if (r > 0.5)
			{
				blobField[i][j] = RED;
			}
			// identify blue pixels
			else if(b > 0.4 && g<b && r<b)
			{
				blobField[i][j] = BLUE;
			}
			// identify green pixels
			else if(g > 0.4 && b<g && r<g)
			{
				blobField[i][j] = GREEN;
			}
		}
	}

	// initialize blob-merging list
	for(int i = 0; i < COLOR_COUNT; i++)
	{
		blobs[i].push_back(-1);
		blobs[i].push_back(-1);
	}

	// now do blob detection
	for(int i = 0; i < height; i++)
	{
		for(int j = 0; j < width; j++)
		{
			// if this is part of a blob
			if(blobField[i][j] != 0)
			{
				// if the top and left are different
				if(i > 0 && j > 0 && blobField[i-1][j] == blobField[i][j-1] 
					&& countField[i-1][j] != countField[i][j-1]
					&& blobField[i-1][j] == blobField[i][j])
				{
					// use the minimum blob number and set them to merge
					countField[i][j] = min(countField[i-1][j], countField[i][j-1]);
					blobs[blobField[i][j]][max(countField[i-1][j], countField[i][j-1])] = min(countField[i-1][j], countField[i][j-1]);
				}
				// check left
				else if(j > 0 && blobField[i][j-1] == blobField[i][j])
					countField[i][j] = countField[i][j-1];
				// check top
				else if(i > 0 && blobField[i-1][j] == blobField[i][j])
					countField[i][j] = countField[i-1][j];
				else
				{
					// add a new blob counter
					int counter = index[blobField[i][j]];
					countField[i][j] = counter;
					blobs[blobField[i][j]].push_back(-1);
					index[blobField[i][j]]++;
				}
			}
		}
	}

	// now do merging

	for(int y = 0; y < height; y++)
	{
		for(int x = 0; x < width; x++)
		{
			int color = blobField[y][x];

			if(color == 0)
				continue;

			if(countField[y][x] >= (int)blobs[color].size())
				cout << "Error " << color << endl;
			if(blobs[color][countField[y][x]] != -1)
			{
				countField[y][x] = blobs[color][countField[y][x]];
				// merge as much as possible
				while(blobs[color][countField[y][x]] != -1)
					countField[y][x] = blobs[color][countField[y][x]];
			}
			// now add to count
			pixelCount[color][countField[y][x]]++;
		}
	}

	// now find largest blobs
	map<int, int>::iterator iter;



	// find largest red blob
	int maxRed = -1;
	for(iter = pixelCount[RED].begin(); iter != pixelCount[RED].end(); iter++)
	{
		if(maxRed == -1)
		{
			maxRed = iter->first;
			continue;
		}

		if(iter->second > pixelCount[RED][maxRed])
			maxRed = iter->first;
	}

	// find largest green blob
	int maxGreen = -1;
	for(iter = pixelCount[GREEN].begin(); iter != pixelCount[GREEN].end(); iter++)
	{
		if(maxGreen == -1)
		{
			maxGreen = iter->first;
			continue;
		}

		if(iter->second > pixelCount[GREEN][maxGreen])
			maxGreen = iter->first;
	}

	// find largest blue blob
	int maxBlue = -1;
	for(iter = pixelCount[BLUE].begin(); iter != pixelCount[BLUE].end(); iter++)
	{
		if(maxBlue == -1)
		{
			maxBlue = iter->first;
			continue;
		}

		if(iter->second > pixelCount[BLUE][maxBlue])
			maxBlue = iter->first;
	}

	// find largest skin blob
	int maxSkin = -1;
	for(iter = pixelCount[SKIN].begin(); iter != pixelCount[SKIN].end(); iter++)
	{
		if(maxSkin == -1)
		{
			maxSkin = iter->first;
			continue;
		}

		if(iter->second > pixelCount[SKIN][maxSkin])
			maxSkin = iter->first;
	}


	// now recolor and process
	for(int y = 0; y < height; y++)
	{
		for(int x = 0; x < width; x++)
		{
			// red
			if(blobField[y][x] == RED && countField[y][x] == maxRed)
			{
				// add to blob
				redBlob.addPixel(x, y);

				if(pixelCount[RED][maxRed] > PIXEL_THRESHOLD)
				{
					// access colors
					uchar* red = ((uchar*)(data + step*y + channels*x + 2));
					uchar* green = ((uchar*)(data + step*y + channels*x + 1));
					uchar* blue = ((uchar*)(data + step*y + channels*x));

					*red = 255;
					*green = 0;
					*blue = 0;
				}
			}

			// green
			else if(blobField[y][x] == GREEN && countField[y][x] == maxGreen)
			{
				// add to blob
				greenBlob.addPixel(x, y);

				if(pixelCount[GREEN][maxGreen] > PIXEL_THRESHOLD)
				{
					// access colors
					uchar* red = ((uchar*)(data + step*y + channels*x + 2));
					uchar* green = ((uchar*)(data + step*y + channels*x + 1));
					uchar* blue = ((uchar*)(data + step*y + channels*x));

					*red = 0;
					*green = 255;
					*blue = 0;
				}
			}

			// blue
			else if(blobField[y][x] == BLUE && countField[y][x] == maxBlue)
			{
				// add to blob
				blueBlob.addPixel(x, y);

				if(pixelCount[BLUE][maxBlue] > PIXEL_THRESHOLD)
				{
					// access colors
					uchar* red = ((uchar*)(data + step*y + channels*x + 2));
					uchar* green = ((uchar*)(data + step*y + channels*x + 1));
					uchar* blue = ((uchar*)(data + step*y + channels*x));

					*red = 0;
					*green = 0;
					*blue = 255;
				}
			}

			// skin/pedal
			else if(blobField[y][x] == SKIN && countField[y][x] == maxSkin)
			{
				// add to blob
				pedalBlob.addPixel(x, y);

				if(pixelCount[SKIN][maxSkin] > 300)
				{
					// access colors
					uchar* red = ((uchar*)(data + step*y + channels*x + 2));
					uchar* green = ((uchar*)(data + step*y + channels*x + 1));
					uchar* blue = ((uchar*)(data + step*y + channels*x));

					*red = 255;
					*green = 255;
					*blue = 0;
				}
			}
		}
	}

}

/*  Calibrate
	Used to perform white-balance calibration (currently unused).
*/
/*void ImageProcessor::Calibrate(IplImage* image)
{
	int width = image->width;
	int height = image->height;
	data = (uchar*)image->imageData;

	for(int i = 0; i < height; i++)
	{
		for(int j = 0; j < width; j++)
		{
			uchar* red = ((uchar*)(data + step*i + channels*j + 2));
			uchar* green = ((uchar*)(data + step*i + channels*j + 1));
			uchar* blue = ((uchar*)(data + step*i + channels*j));

			*red += (int)conversionR;
			*green += (int)conversionG;
			*blue += (int)conversionB;

			/*int average = (*red + *blue + *green)/3;

			if(average > 128)
			{
				*red -= (uchar)((average - 128.0)/2);
				*green -= (uchar)((average - 128.0)/2);
				*blue -= (uchar)((average - 128.0)/2);
			}
			else
			{
				// make lighter
				*red += (uchar)((128.0 - average)/255.0)**red;
				*green -= (uchar)((128.0 - average)/255.0)**green;
				*blue -= (uchar)((128.0 - average)/255.0)**blue;
			}
		}
	}
}
*/